#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import pytyrant


class TableError(Exception):
  # error code: success
  ESUCCESS = 0
  # error code: invalid operation
  EINVALID = 1
  # error code: host not found
  ENOHOST = 2
  # error code: connection refused
  EREFUSED = 3
  # error code: send error
  ESEND = 4
  # error code: recv error
  ERECV = 5
  # error code: existing record
  EKEEP = 6
  # error code: no record found
  ENOREC = 7
  # error code: miscellaneous error
  EMISC = 9999

  def __str__(self):
    code = int(self.args[0])
    if code == self.ESUCCESS:
      return 'success'
    elif code == self.EINVALID:
      return 'invalid operation'
    elif code == self.ENOHOST:
      return 'host not found'
    elif code == self.EREFUSED:
      return 'connection refused'
    elif code == self.ESEND:
      return 'send error'
    elif code == self.ERECV:
      return 'recv error'
    elif code == self.EKEEP:
      return 'existing record'
    elif code == self.ENOREC:
      return 'no record found'
    elif code == self.EMISC:
      return 'miscellaneous error'
    return 'unknown'


class IndexType(object):
  # index type: lexical string
  LEXICAL = 0
  # index type: decimal string
  DECIMAL = 1
  # index type: optimize
  # use this after using LEXICAL or DECIMAL
  OPT = 9998
  # index type: void
  VOID = 9999
  # index type: keep existing index
  KEEP = 1 << 24


class Table(pytyrant.Tyrant):
  '''Helper function.'''
  def _put(self, key, cols, command):
    if not isinstance(cols, dict):
      raise TypeError('cols must be a dictionary.')
    args = [key]
    for item in cols.items():
      args.extend(item)
    return bool(self.misc(command, 0, args))

  def put(self, key, cols):
    '''Store a record.
    `pkey` specifies the primary key.
    `cols` specifies a hash containing columns.
    If a record with the same key exists in the database, it is overwritten.'''
    return self._put(key, cols, 'put')

  def putkeep(self, key, cols):
    '''Store a new record.
    `pkey` specifies the primary key.
    `cols` specifies a hash containing columns.
    If a record with the same key exists in the database, this method has no
    effect.'''
    return self._put(key, cols, 'putkeep')

  def putcat(self, key, cols):
    '''Concatenate columns of the existing record.
    `pkey` specifies the primary key.
    `cols` specifies a hash containing columns.
    If there is no corresponding record, a new record is created.'''
    return self._put(key, cols, 'putcat')

  def get(self, key):
    '''Retrieve a record.
    `pkey` specifies the primary key.
    The return value is a dictionary of the columns of the corresponding record.'''
    result = self.misc('get', 0, [key])
    keys = result[::2]
    values = result[1::2]
    return dict(zip(keys, values))

  def mget(self, keys):
    '''Retrieve records.
    `keys` specifies a list containing the retrieval keys.
    Due to the protocol restriction, this method can not handle records with binary
    columns including the "\0" character.'''
    results = super(Table, self).mget(keys)
    new_results = {}
    for k, v in results:
      keys = v.split('\0')[::2]
      values = v.split('\0')[1::2]
      new_results[k] = dict(zip(keys, values))
    return new_results

  def setindex(self, name, type):
    '''Set a column index.
    `name` specifies the name of a column. If the name of an existing index is
    specified, the index is rebuilt. An empty string means the primary key.
    `type` specifies the index type (IndexType).'''
    return bool(self.misc('setindex', 0, [name, str(type)]))

  def genuid(self):
    '''Generate a unique ID number. The return value is the new unique ID
    number or None on failure.'''
    result = self.misc('genuid', 0, [])
    if result:
      return result[0]
    return None

  def misc(self, func, opts, args):
    try:
      return super(Table, self).misc(func, opts, args)
    except pytyrant.TyrantError, e:
      raise TableError(*e.args)


class QueryCondition(object):
  # query condition: string is equal to
  STREQ = 0
  # query condition: string is included in
  STRINC = 1
  # query condition: string begins with
  STRBW = 2
  # query condition: string ends with
  STREW = 3
  # query condition: string includes all tokens in
  STRAND = 4
  # query condition: string includes at least one token in
  STROR = 5
  # query condition: string is equal to at least one token in
  STROREQ = 6
  # query condition: string matches regular expressions of
  STRRX = 7
  # query condition: number is equal to
  NUMEQ = 8
  # query condition: number is greater than
  NUMGT = 9
  # query condition: number is greater than or equal to
  NUMGE = 10
  # query condition: number is less than
  NUMLT = 11
  # query condition: number is less than or equal to
  NUMLE = 12
  # query condition: number is between two tokens of
  NUMBT = 13
  # query condition: number is equal to at least one token in
  NUMOREQ = 14
  # query condition: negation flag
  NEGATE = 1 << 24
  # query condition: no index flag
  NOIDX = 1 << 25
  # order type: string ascending


class QueryOrder(object):
  STRASC = 0
  # order type: string descending
  STRDESC = 1
  # order type: number ascending
  NUMASC = 2
  # order type: number descending
  NUMDESC = 3


class Query(object):
  def __init__(self, rdb=None):
    '''Create a query object. `rdb` specifies the remote database object.
    The return value is the new query object.'''
    self.rdb = rdb
    self.args = []

  def add_cond(self, name, op, expr):
    '''Add a narrowing condition.
    `name` specifies the name of a column. An empty string means the primary
    key.
    `op` specifies an operation type (QueryCondition).
    `expr` specifies an operand expression.'''
    self.args.append('addcond\0%s\0%s\0%s' % (name, op, expr))

  def set_order(self, name, type):
    '''Set the order of the result.
    `name` specifies the name of a column. An empty string means the primary
    key.
    `type` specifies the order type (QueryOrder).'''
    self.args.append('setorder\0%s\0%s' % (name, type))

  def set_limit(self, max=-1, skip=-1):
    '''Set the maximum number of records of the result.
    `max` specifies the maximum number of records of the result. If it is not
    defined or negative, no limit is specified.
    `skip` specifies the maximum number of records of the result. If it is not
    defined or not more than 0, no record is skipped.'''
    self.args.append('setlimit\0%s\0%s' % (max, skip))

  def search(self):
    '''Execute the search. The return value is a list of the primary keys of
    the corresponding records. This method never fails and returns an empty
    list even if no record corresponds.'''
    return self.rdb.misc('search', pytyrant.RDBMONOULOG, self.args)

  def search_out(self):
    '''Remove each corresponding record.'''
    return bool(self.rdb.misc('search', 0, self.args + ['out']))

  def search_get(self, names=None):
    '''Get records corresponding to the search.
    `names` specifies an array of column names to be fetched. An empty string
    means the primary key. If it is not defined, every column is fetched.
    The return value is an array of column hashes of the corresponding records.
    This method does never fail and return an empty list even if no record
    corresponds.
    Due to the protocol restriction, this method can not handle records with
    binary columns including the "\0" chracter.'''
    if names is None:
      args = self.args + ['get']
    else:
      args = self.args + ['get\0\0%s' % '\0'.join(names)]
    results = self.rdb.misc('search', 0, args)
    new_results = {}
    for result in results:
      result = result.split('\0')
      keys = result[2::2]
      values = result[3::2]
      new_results[result[1]] = dict(zip(keys, values))
    return new_results


if __name__ == '__main__':
  # Some examples on how to use the library.
  t = Table.open()
  for i in range(1, 10):
   t.put(str(i), {'id': str(i)})
  t.putcat('1', {'foo': 'bar'})
  print t.get('1')
  print t.mget(['2', '3', '4', '5'])

  q = Query(t)
  q.add_cond('foo', QueryCondition.STREQ, 'bar')
  print q.search_get()
  print q.search_get(names=['foo'])

  q = Query(t)
  q.add_cond('id', QueryCondition.NUMGE, '3')
  q.set_limit(4)
  print q.search()
  q.set_order('id', QueryOrder.NUMDESC)
  print q.search()
